import sys

import numpy as np 
sys.path.append('..')

from common.layers import Embedding


class EmbeddingDot:
    def __init__(self,W):
        self.embed = Embedding(W)
        self.params = self.embed.params
        self.grads = self.embed.grads 
        self.cache = None 

    def forward(self,h,idx):
        target_W = self.embed.forward(idx)
        out = np.sum(target_W*h, axis=1)
        
        self.cache = (h,target_W)
        return out 
    
    def backward(self,dout):
        h,target_W = self.cache
        dout = dout.reshape(dout.shape[0],1)

        dtarget_W = dout*h
        self.embed.backward(dtarget_W)
        dh = dout*target_W
        return dh 